import torch
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import ImageFolder
from torchvision.transforms import transforms


def generate_cars_loader(train_path="./data/train", valid_path="./data/val", batch_size=10, input_size=(224, 224),
                        in_channels=3):
    # transform 不会进行数据持久化（通俗，不会永久改变数据集）
    # 每当epoch训练时，transform都会随机数据增强进行训练。（泛化性会大大提高）
    transform = transforms.Compose([
        transforms.RandomRotation((-10, 10)),
        transforms.Resize(input_size),  # 修改原图像的大小（线性差值）
        transforms.ToTensor()  # 将PIL格式转化为torch.tensor(张量 --> 矩阵)
    ])
    if in_channels == 1:  # 如果要求是一个通道，则输入必须进行灰度化处理
        transform = transforms.Compose([
            transforms.Resize(input_size),
            transforms.RandomRotation((-10, 10)),
            transforms.Grayscale(),  # 进行灰度化处理
            transforms.ToTensor()
        ])
    train_dataset = ImageFolder(train_path, transform=transform)
    valid_dataset = ImageFolder(valid_path, transform=transform)
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
    valid_loader = DataLoader(valid_dataset, batch_size=batch_size)
    return train_loader, valid_loader
